# 从csv文件中读取数据，分别为：X列表和对应的Y列表

import pandas as pd
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression


def get_data():
    # 加载CSV文件
    df = pd.read_csv('./linear_regression/price_info.csv')

    # 提取特征列和标签列
    X = df['square_feet'].values.reshape(-1, 1)  # 转换为二维数组
    Y = df['price'].values  # 目标向量（价格）

    return X, Y


# 获取线性回归模型
def get_linear_model(X, Y):
    # 1.创建线性回归模型
    model = LinearRegression()

    # 2.使用X/Y对模型进行拟合
    model.fit(X, Y)

    return model


def show_linear_regression(X, Y, model):
    # 绘制原始数据点
    plt.scatter(X, Y, color='blue', label='Actual Price')

    # 绘制回归线
    plt.plot(X, model.predict(X), color='red', linewidth=2, label='Prediction Line')

    # 图表装饰
    plt.title('House Price Prediction')
    plt.xlabel('Square Feet')
    plt.ylabel('Price')
    plt.legend()

    # 显示图像
    plt.show()


def main():
    # 1. 读取数据
    X, Y = get_data()

    # 2. 模型训练
    model = get_linear_model(X, Y)

    # 3. 获取700平方英尺时的预测价格
    prediction = model.predict([[700]])
    print(f'700平方英尺的预测价格为：${prediction[0]:.2f}')

    # 4. 绘图
    show_linear_regression(X, Y, model)


if __name__ == '__main__':
    main()